import numpy as np
import os
from tqdm import trange
from ..numbajit import *


class moriAnalysis(object):

    def __init__(self, **kwargs):
        self.wd = kwargs.get('wd', os.getcwd())
        self.filename = kwargs.get('filename', None)
        self.fmt = kwargs.get('fmt', 'csv')
        self.ori = kwargs.get('ori', None)
        self.numIter = self.ori.shape[0]

    def calcKAM(self, **kwargs):

        workList = kwargs.get('worklist', None)
        wlt = kwargs.get('wltype', 1)

        if wlt == 1:
            if not isinstance(workList, list):
                workList = workList.readlines()
        else:
            pass

        kamList = np.array(())
        for i in trange(self.numIter):
            tempKAM = 0
            if wlt == 1:
                tempWorkList = np.fromstring(workList[i], dtype=int, sep='\t')
            else:
                tempWorkList = workList[i]

            numNeigh = len(tempWorkList)
            for j in tempWorkList:
                q1 = self.ori[i]
                q2 = self.ori[j - 1]
                quatSym1 = symmetriesCubic(q1)
                quatSym2 = symmetriesCubic(q2)
                mang = calcDisoriCubic(quatSym1, quatSym2)
                tempKAM += mang
            tempKAM = tempKAM / numNeigh
            kamList = np.append(kamList, tempKAM)

        np.savetxt(os.path.join(self.cwd, '{}_KAM.{}'.format(self.filename, self.fmt)), kamList, fmt='%1.3f',
                   delimiter='\t')

    def calcGAM(self, **kwargs):

        workList = kwargs.get('worklist', None)
        if not isinstance(workList, list):
            workList = workList.readlines()

        gamList = np.array(())
        for i in trange(self.numIter):
            tempGAM = 0
            tempWorkList = np.fromstring(workList[i], dtype=int, sep='\t')
            numNeigh = len(workList)

            if tempWorkList.any() == 0:
                tempGAM = np.array([0])
            else:
                for j in tempWorkList:
                    q1 = self.ori[i]
                    q2 = self.ori[j - 1]
                    quatSym1 = symmetriesCubic(q1)
                    quatSym2 = symmetriesCubic(q2)
                    mang = calcDisoriCubic(quatSym1, quatSym2)
                    tempGAM += mang

                tempGAM = tempGAM / numNeigh

            gamList = np.append(gamList, tempGAM)

        np.savetxt(os.path.join(self.wd, '{}_GAM.{}'.format(self.filename, self.fmt)), gamList, fmt='%1.3f',
                   delimiter='\t')

    def calcGOS(self, **kwargs):

        elesetArray = kwargs.get('eleset', None)
        wl = kwargs.get('wl', False)

        if wl:
            workList = kwargs.get('worklist', None)
            if not isinstance(workList, list):
                workList = workList.readlines()

        elesetNum = np.unique(elesetArray)
        numEleset = elesetNum.shape[0]

        gosList = np.zeros(self.numIter)
        for i in trange(numEleset):
            tempGOS = 0

            if not wl:
                tempWorkList = np.where(elesetArray == elesetNum[i])[0]
            else:
                tempWorkList = np.fromstring(workList[i], dtype=int, sep='\t')

            numNeigh = len(tempWorkList)
            tempQuatArray = self.ori[tempWorkList]
            q2 = np.mean(tempQuatArray, axis=0)
            quatSym2 = symmetriesCubic(q2)

            for j in tempWorkList:
                q1 = self.ori[j]
                quatSym1 = symmetriesCubic(q1)
                mang = calcDisoriCubic(quatSym1, quatSym2)
                tempGOS += mang

            tempGOS = tempGOS / numNeigh
            gosList[workList] = tempGOS

        np.savetxt(os.path.join(self.wd, '{}_GOS.{}'.format(self.filename, self.fmt)), gosList, fmt='%1.3f',
                   delimiter='\t')